from data_generation import Long_Data_gener
import numpy as np
from scipy import special
from nonparameter_classfication import NonparaClassfication
import parameter_classfication
import pandas as pd
from tqdm import tqdm

import pyjags
def dina_parameter_estimate(N,Q,I,K,simu_resp,chains=2,cores=-1,burning=5000,iteration=5000,return_trace=False,return_full_vars=False):
    # hodina_pyjags
    jags_ho_dina_data = {
        "Y": simu_resp,
        "Q": Q,
        "N": N,
        "K": K,
        "I": I

    }
    if return_full_vars == True:
        jags_model_return_vars = ["p", "alpha_p", "alpha", "theta", "beta", "lambda_", "s", "g"]
    else :
        jags_model_return_vars = ["alpha"]

    jags_hodina_model_str =\
    """
    model{
        for (n in 1:N) {
            for (i in 1:I){
                for (k in 1:K){
                    w[n,i,k] <-(alpha[n,k]>=Q[i,k])
                }
                eta[n,i] <- prod(w[n,i,1:K])
                p[n,i] <- eta[n,i]*(1-s[i]-g[i]) + g[i]
                Y[n,i] ~ dbern(p[n,i])
            }
        }

        for (n in 1:N){
            for (k in 1:K){
                logit(alpha_p[n,k]) <- lambda_[k]*theta[n]+beta[k]
                alpha[n,k] ~ dbern(alpha_p[n,k]) 
            }
            theta[n] ~ dnorm(0,1)
        }

        for (k in 1:K){
            beta[k] ~ dnorm(0,0.25)
            lambda_[k] ~ dnorm(0,0.25) T(0,)
        } 
        for (i in 1:I){
            s[i] ~ dbeta(1,1)
            g[i] ~ dbeta(1,1) T(,1-s[i])
        }
    }"""
    jags_ho_dina_model = pyjags.Model(code=jags_hodina_model_str,
                                        data=jags_ho_dina_data,
                                        chains=chains,
                                        threads=cores,
                                        generate_data=True,
                                        chains_per_thread=1,
                                        adapt=False,
                                      progress_bar=True
                                        )
    print("burning.......")
    jags_tarce_prior = jags_ho_dina_model.sample(iterations=burning, vars=[], thin=1)
    print("sampling.......")
    jags_tarce_posterior = jags_ho_dina_model.sample(iterations=iteration, vars=jags_model_return_vars, thin=1)
    x = (jags_tarce_posterior["alpha"].mean(axis=-1).mean(axis=-1) > 0.5)
    if return_trace==True:
        return x,jags_tarce_posterior
    return x


def long_dina_parameter_estimate(T,N,Q,I,K,simu_resp,burning=5000,iteration=5000,chains=2,cores=-1
                                 ,return_trace=False,return_full_vars=False):
    jags_data = {
        "Y": simu_resp,
        "Q": Q,
        "N": N,
        "I": I,
        "K": K,
        "T": T
    }
    if return_full_vars==True:
        jags_return_vars = ["mu_theta", "pr_theta", "theta", "att_possi",
                        "att", "theta_beta", "theta_lambda", "lambda_0", "lambda_k", "p"]
    else :
        jags_return_vars = ["att","att_possi"]

    long_dina_jags_model = ling_dina_jags_model = \
        """
        model{
            for (t in 1:T){
                for (n in 1:N){
                    for (i in 1:I){
                        for (k in 1:K){w[t,n,i,k] <- (att[t,n,k]>=Q[t,i,k])}
                        eta[t,n,i] <- prod(w[t,n,i,1:K])
                        logit(p[t,n,i]) <- lambda_0[t,i]+lambda_k[t,i]*eta[t,n,i]
                        Y[t,n,i] ~ dbern(p[t,n,i])
                    }
                }
    
    
                for (n in 1:N){
                    for (k in 1:K){
                        logit(att_possi[t,n,k]) <- theta_beta[k] + theta_lambda[k]*theta[t,n]
                        att[t,n,k] ~ dbern(att_possi[t,n,k])
                    }
                }
            }
    
            for (n in 1:N){
                theta[1:T,n] ~ dmnorm(mu_theta[1:T], pr_theta[1:T,1:T])
            }
            for (k in 1:K){
                theta_beta[k]~ dnorm(0,0.25)
                theta_lambda[k] ~ dnorm(0,0.25) T (1,)
            }
    
            for (t in 1:T){
                for (i in 1:I){
                    lambda_0[t,i] ~ dnorm(-1.096, 0.25)
                    lambda_k[t,i] ~ dnorm(0, 0.25) T(0, )
                }
            }
    
            mu_theta[1] <- 0
            for (t in 2:T){
                mu_theta[t] ~ dnorm(0,0.5)
            }
    
            L_theta[1,1] <- 1
            for (tt in 2:T){
                L_theta[tt,tt] ~ dgamma(1,1)
                for (ttt in 1:(tt-1)){
                    L_theta[tt,ttt] ~ dnorm(0,1)
                    L_theta[ttt,tt] <- 0
                }
            }
            sigma_theta = L_theta %*% t(L_theta)
            pr_theta[1:T, 1:T] <- inverse(sigma_theta[1:T, 1:T])
    
    
        }
        """
    jags_long_dina_model = pyjags.Model(code=long_dina_jags_model,
                                        data=jags_data,
                                        chains=chains,
                                        threads=cores,
                                        generate_data=True,
                                        chains_per_thread=1,
                                        adapt=False,
                                        progress_bar=True
                                        )
    print("burning.......")
    jags_tarce_prior = jags_long_dina_model.sample(iterations=burning, vars=[], thin=1)
    print("sampling.......")
    jags_tarce_posterior = jags_long_dina_model.sample(iterations=iteration, vars=jags_return_vars, thin=1)

    x = (jags_tarce_posterior["att"].mean(axis=-1).mean(axis=-1) > 0.5)
    if return_trace==True:
        return x,jags_tarce_prior
    return x





def get_indicator(att_pre, att_real):
    x = att_pre == att_real
    ACCR = np.nanmean(x)
    y = np.product(x, axis=-1)
    PCCR = np.nanmean(np.product(x, axis=-1))
    Long_PCCR = np.nanmean(np.product(np.product(x, axis=-1), axis=0))
    return ACCR, PCCR, Long_PCCR


def run(t, k, i, N, s, g):
    theta_beta = np.ones((k)) * 1.5
    theta_lambda = np.ones((k)) * (-1) + np.array([i * 0.5 for i in range(k)])
    lambda_0 = np.repeat(special.logit(g), i * t).reshape(t, i)
    lambda_k = np.repeat(special.logit(1 - s) - special.logit(g), i * t).reshape(t, i)
    # 实例化对象
    gener = Long_Data_gener()
    classf = NonparaClassfication()
    # 生成被试高阶能力
    theta_cov = np.array([[1, 1.25, 1.76, 2.74, 4.26],
                         [1.25, 1.56, 1.76, 2.19, 2.73],
                         [1.76, 1.76, 2.43, 3.8, 4.26],
                         [2.74, 2.19, 2.73, 3.8, 4.26],
                         [4.26, 2.73, 3.41, 4.26, 5.9]])
    if t == 5:
        theta = gener.m_theta(t, mean=0, growth_rate=0.5, n=N, input_cor_mat=theta_cov)
    elif t==3:
        theta = gener.m_theta(t, mean=0, growth_rate=0.5, n=N, input_cor_mat=np.array(theta_cov[:3,:3]))
    elif t==2:
        theta = gener.m_theta(t, mean=0, growth_rate=0.5, n=N, input_cor_mat=np.array(theta_cov[:2,:2]))

    # 生成被试属性
    att, z = gener.attribute(theta, theta_beta, theta_lambda, return_possi=True)
    # 生成题目Q矩阵
    Q = gener.Q_generate_long(t, i, k)
    # 生成模拟作答
    resp, resp_possi = gener.simu_resp(t, att, Q, lambda_k=lambda_k, lambda_0=lambda_0, return_possi=True)

    """被试数据生成到此结束"""
    # 获取被试所有可能的属性掌握模式
    all_pattern = gener.all_pattern(k)
    # 计算所有可能的属性掌握模式对应的理想作答
    class_ideal_resp = gener.long_dina_ideal_resp(t, Q, all_pattern)
    # 参数方法
    para_pre_att = parameter_classfication.dina_parameter_estimate(N,Q[0],I=i,K=k,simu_resp=resp[0],return_trace=False)
    # para_pre_att = att[0]
    # para_pre_att = att[0]
    # 使用非参方法对被试进行分类
    class_result_dic = {}
    class_result_dic["class_result1"] = classf.long_standard_distance(t, all_pattern, simu_resp=resp,
                                                                      class_ideal_response=class_ideal_resp,
                                                                      distance="hamming")
    class_result_dic["class_result2"], l = classf.long_standard_dis_constrain_1(t, all_pattern, simu_resp=resp,
                                                                                class_ideal_response=class_ideal_resp,
                                                                                distance="hamming", return_judge=True)
    class_result_dic["class_result3"] = classf.long_standard_dis_constrain_2(t, all_pattern, simu_resp=resp,
                                                                             class_ideal_response=class_ideal_resp,
                                                                             distance="hamming")
    class_result_dic["class_result4"], l = classf.long_standard_dis_constrain_1(t, all_pattern, simu_resp=resp,
                                                                                class_ideal_response=class_ideal_resp,
                                                                                distance="hamming", return_judge=True,
                                                                                input_first_t_att=att[0, :, :])
    class_result_dic["class_result5"] = classf.long_standard_dis_constrain_2(t, all_pattern, simu_resp=resp,
                                                                             class_ideal_response=class_ideal_resp,
                                                                             distance="hamming",
                                                                             input_first_t_att=att[0, :, :])
    class_result_dic["class_result6"] = classf.long_standard_distance(t, all_pattern, simu_resp=resp,
                                                                      class_ideal_response=class_ideal_resp,
                                                                      distance="hamming", weight=True)
    class_result_dic["class_result7"], l = classf.long_standard_dis_constrain_1(t, all_pattern, simu_resp=resp,
                                                                                class_ideal_response=class_ideal_resp,
                                                                                distance="hamming", return_judge=True,
                                                                                input_first_t_att=para_pre_att)
    class_result_dic["class_result8"] = classf.long_standard_dis_constrain_2(t, all_pattern, simu_resp=resp,
                                                                             class_ideal_response=class_ideal_resp,
                                                                             distance="hamming",
                                                                             input_first_t_att=para_pre_att)
    class_result_dic["class_result9"] = classf.long_standard_dis_constrain_2(t, all_pattern, simu_resp=resp,
                                                                             class_ideal_response=class_ideal_resp,
                                                                             distance="weight_hamming",
                                                                             input_first_t_att=att[0, :, :])
    class_result_dic["class_result10"] = classf.long_standard_dis_constrain_2(t, all_pattern, simu_resp=resp,
                                                                              class_ideal_response=class_ideal_resp,
                                                                              distance="weight_hamming",
                                                                              input_first_t_att=para_pre_att)
    # class_result_dic["class_result11"] = classf.long_standard_dis_constrain_2(t, all_pattern, simu_resp=resp,
    #                                                                           class_ideal_response=class_ideal_resp,
    #                                                                           distance="weight_hamming",
    #                                                                           input_first_t_att=para_pre_att)
    class_result_dic["class_result11"] = parameter_classfication.long_dina_parameter_estimate(T=t,N=N,Q=Q,I=i,K=k,simu_resp=resp)

    class_mathod_name_list = ["标准", "限制1", '限制2',
                              '第一次完全正确限制1', '第一次完全正确限制1',
                              "加权距离", "第一次为参数估计限制1", "第一次为参数估计限制2",
                              "第一次完全正确加权距离限制2", "第一次为参数估计加权距离限制2", "Long_dina_参数"]
    rename_dict = {}
    for i in range(len(class_mathod_name_list)):
        att_pre = class_result_dic["class_result{}".format(i + 1)]
        A, P, L = get_indicator(att_pre, att)
        rename_dict["ACCR_" + class_mathod_name_list[i]] = A
        rename_dict["PCCR_" + class_mathod_name_list[i]] = P
        rename_dict["Long_PCCR_" + class_mathod_name_list[i]] = L

    return rename_dict


R = 5 
I_list = [50, 30]
T_list = [3，2，5]
K_list = [5,3]
N_list = [300, 100]

pbar = tqdm(total=len(I_list) * len(T_list) * len(K_list) * len(N_list) * R)
for t in T_list:
    for k in K_list:
        for N in N_list:
            for i in I_list:
                result_list = []
                for r in range(R):
                    rename_dict = run(t=t,
                                      k=k,
                                      i=i,
                                      N=N,
                                      s=0.1,
                                      g=0.1)
                    result_list.append(rename_dict)
                    pbar.update(1)

                df = pd.DataFrame(result_list)
                df.to_csv("./result1/result_t{}_k{}_n{}_i{}.csv".format(t, k, N, i), index=False)
